package com.leetcode.No1337;

import org.junit.Test;

import java.util.Arrays;
import java.util.Comparator;
import java.util.TreeSet;

/**
 * @program: leetcode
 * @description: 矩阵中战斗力最弱的 K 行
 * @author: wangzhihua
 * @date: 2021-08-05
 */
public class Solution01 {

	@Test
	public void test01() {
		int[][] arrays = new int[][] {
				{1,1,0,0,0},
				{1,1,1,1,0},
				{1,0,0,0,0},
				{1,1,0,0,0},
				{1,1,1,1,1}};
		System.out.println(Arrays.toString(kWeakestRows(arrays, 4)));
	}

	@Test
	public void test02() {
		int[][] arrays = new int[][] {
				{1,0,0,0},
				{1,1,1,1},
				{1,0,0,0},
				{1,0,0,0}};
		System.out.println(Arrays.toString(kWeakestRows(arrays, 4)));
	}

	public int[] kWeakestRows(int[][] mat, int k) {
		NumInfo[] tmpAns = new NumInfo[mat.length];
		for (int i = 0; i < mat.length; ++i) {
			int sum = 0;
			for (int num : mat[i]) {
				if (num == 0) {
					break;
				}
				sum += num;
			}
			tmpAns[i] = new NumInfo(i, sum);
		}
		Arrays.sort(tmpAns, new Comparator<NumInfo>() {
			@Override
			public int compare(NumInfo o1, NumInfo o2) {
				return (o1.val == o2.val) ? (o1.col - o2.col) : (o1.val - o2.val);
			}
		});

		int[] ans = new int[k];
		for (int i = 0; i < k; ++i) {
			ans[i] = tmpAns[i].col;
		}

		return ans;
	}

	private class NumInfo {
		int col;
		int val;
		NumInfo(int col, int val) {
			this.col = col;
			this.val = val;
		}
	}

}
